package com.gmrz.uaf.crypto.internal.handlers;

import com.gmrz.uaf.common.UAFRuntimeException;
import com.gmrz.uaf.crypto.CryptoConstants;
import org.apache.commons.codec.DecoderException;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.codec.binary.Hex;
import org.apache.commons.lang3.ArrayUtils;
import org.bouncycastle.jcajce.provider.symmetric.AES;
import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.GeneralSecurityException;
import java.security.SecureRandom;
import java.security.Security;
import java.util.Arrays;
import java.util.Random;

/**
 * Created by chentaikun on 17/1/18.
 */
public class ServerDataKeyParser {
    private final String wrapKey = "2ca1601e7f764291b8dfeb0357bdd5b9";
    private final String keyAlgorithm = "AES";
    private final String transformation = "AES/GCM/NoPadding";

    public SecretKey getKey()
    {
        try {
            return new SecretKeySpec(Hex.decodeHex(wrapKey.toCharArray()), this.keyAlgorithm);
        } catch (DecoderException e) {
            e.printStackTrace();
        }
        return null;
    }
    public byte[] encrypt(byte[] clearData) throws GeneralSecurityException {
        byte[] iv = null;
        IvParameterSpec ivSpec = null;
        Cipher cipher = Cipher.getInstance(this.transformation, CryptoConstants.BOUNCY_CASTLE_PROVIDER_NAME);

        iv = generateIV(cipher.getBlockSize());
        ivSpec = new IvParameterSpec(iv);
        cipher.init(1, getKey(), ivSpec);
        return ArrayUtils.addAll(iv, cipher.doFinal(clearData));
    }

    public byte[] decrypt(byte[] cipherText) throws GeneralSecurityException {
        Cipher cipher = null;
        byte[] iv = null;
        IvParameterSpec ivSpec = null;
        cipher = Cipher.getInstance(this.transformation, CryptoConstants.BOUNCY_CASTLE_PROVIDER_NAME);

        if (cipherText == null) {
            throw new UAFRuntimeException("Cipher text is empty.");
        }

        if (cipherText.length < cipher.getBlockSize()) {
            throw new UAFRuntimeException("Cipher text is not of the desired length.");
        }
        iv = Arrays.copyOfRange(cipherText, 0, cipher.getBlockSize());
        ivSpec = new IvParameterSpec(iv);

        byte[] cipherTextMinusIv = Arrays.copyOfRange(cipherText, iv.length, cipherText.length);

        cipher.init(2, getKey(), ivSpec);

        return cipher.doFinal(cipherTextMinusIv);
    }

    byte[] generateIV(int blockSize) {
        byte[] ivBytes = new byte[blockSize];
        SecureRandom random = new SecureRandom();
        random.nextBytes(ivBytes);
        return ivBytes;
    }

    public byte[] generate() throws GeneralSecurityException {
        KeyGenerator keyGenerator = KeyGenerator.getInstance("AES");
        keyGenerator.init(128);
        SecretKey key = keyGenerator.generateKey();
        return key.getEncoded();

    }
    public static void main(String args[]){

        Security.addProvider(new BouncyCastleProvider());
        ServerDataKeyParser parser = new ServerDataKeyParser();

        try {
            byte [] key = parser.generate();

            byte[] wrappedKey = parser.encrypt(key);

            byte[] plainkey = parser.decrypt(wrappedKey);

        } catch (GeneralSecurityException e) {
            e.printStackTrace();
        }
    }

}
